Server: fused single-GPU decode for --batched-session (co-decode 2-4 sessions in one graph pass) - #604
Open
iCreil wants to merge 10 commits into
Open
Server: fused single-GPU decode for --batched-session (co-decode 2-4 sessions in one graph pass)#604iCreil wants to merge 10 commits into
iCreil wants to merge 10 commits into
Conversation
Add ds4_gpu_attention_decode_heads_multi_tensor(): one query token per sequence, each attending to its own raw + compressed caches with decode (everything visible) semantics. Sequence b reads q row b and writes heads row b; per-sequence cache pointers, ring geometry and optional indexer mask row travel in a ds4_gpu_attn_seqview array (max 4, passed by value in the kernel parameters). The kernel is a standalone adaptation of the n_tokens==1 launch of attention_decode_mixed_kernel, so the single-session hot path is untouched. Sequences whose compressed row count exceeds the shared score buffer are rejected for now: the large-context online variant is a follow-up, and callers gate batched decode on this. tests/attention_multi_smoke drives two sequences with different shapes, ring offsets and masks and checks the multi launch against two independent single-sequence launches: bitwise identical on both the head_dim==512 fast path and the generic path (RTX PRO 6000, CUDA 13).
Batched decode evaluates one token per sequence, each at its own arbitrary position, while rope_tail_kernel and head_rms_norm_rope_tail_kernel derive the position as pos0 + t (* stride). Add _multi twins that take a small by-value position array (max DS4_GPU_DECODE_MULTI_MAX entries) indexed by token row, with per-element arithmetic kept identical so each row matches a single-token launch bitwise. tests/rope_multi_bench checks that equivalence (plain and YaRN paths) and measures why the variants exist: on an RTX PRO 6000, per-sequence loops of single-token launches over the 4 decode RoPE sites cost ~1.0 ms extra per 43-layer step at B=4 (~516 launches x ~2 us) versus batched launches.
The batch warp kernel puts n_tok on the grid Y axis, so every token re-reads the same weight row from VRAM; in decode, where weights are the whole cost, that scales the time with n_tok for no benefit. The generic per-row kernel (used when blocks>32) is worse still. The new narrow kernel keeps one warp per output row and, for each 32-weight block it reads once, dots it against all n_tok activation blocks, so the weight bytes cross memory once and are reused from L1. Output layout matches the batch kernel; the n_tok==1 and prefill paths are untouched. Dispatched for n_tok in [2, DS4_CUDA_NARROW_MAX] ahead of cuBLAS and the batch warp kernel; DS4_CUDA_NO_Q8_NARROW forces the old path. Modeled on the n_tok==1 warp8 kernel, it loops blocks with a lane stride and handles any in_dim. Decode-step bench (RTX PRO 6000, V4 Flash q2), aggregate vs interleaving: B=2 1.47x, B=3 1.67x, B=4 1.83x (was B=2 0.62x through the batch path). Correctness: batched_decode_smoke shows narrow and batch-warp produce identical token streams. Batched decode still differs slightly from single decode, but that is the existing non-deterministic MoE reduction (atomicAdd over experts), not the kernel: verified by first-step logits being bit-identical between narrow and batch-warp.
…ode driver) Port of the fork's batched-decode driver onto the current main: - metal_graph_decode_compressor_indexer: per-sequence compressor/indexer helper re-extracted from the current decode layer phase (kept as a copy, single-token path untouched). - metal_graph_encode_layer_attention_decode_multi + encode_decode_multi: dense projections batched on the aliased batch workspace with n_tokens = n_seqs (narrow q8_0 kernel serves 2-4 rows), pos-vector RoPE, per-sequence KV store and compressor/indexer, multi-sequence attention with per-sequence fallbacks, FFN and output head via the existing batch functions. - ds4_sessions_eval_batch_cuda: single-tier decode batches take the fused pass instead of one-token pipelining (DS4_NO_FUSED_SESSION_BATCH=1 restores it); logits land one row per sequence in spec_logits, now allocated unless the fused path is disabled.
In the fused session-batch decode each FFN batch row is one token of a distinct sequence, so the selected experts are almost always disjoint and the generic batch expert path costs more than n_rows single-token passes. Route the routed-MoE dispatch of metal_graph_encode_layer_ffn_batch through ds4_gpu_routed_moe_one_tensor per row while the fused driver is on the stack (file-scope flag; the graph worker serializes encodes). DS4_FUSED_BATCH_MOE=1 keeps the batch expert path for comparison.
With --batched-session N the lazy Q8->F16 weight caches warm up during engine creation and grow until only the default reserve (~1% of VRAM) stays free, so the resident sessions allocated afterwards fail even at modest context sizes (on a 96 GB card even one 32k session could not be created). Derive the reserve from the per-session context-buffer estimate times N plus the shared prefill workspace, and export it before engine creation; an explicit DS4_CUDA_Q8_F16_CACHE_RESERVE_MB wins.
Move the reserve computation into engine open, right before the startup cache warm: the context-memory estimate needs the model shape, which is only loaded here (the earlier server-side attempt underestimated it and still failed session creation). The server passes the resident session count via the new resident_sessions config field.
This was referenced Jul 25, 2026
iSevenDays
added a commit
to iSevenDays/ds4
that referenced
this pull request
Jul 25, 2026
…tirez#605 race fix Cherry-picked onto feat/mtp-native (which has upstream main + MTP/DSpark + expert LRU): - antirez#604: 10 commits (session-batch decode, fused routed experts, VRAM reserve for resident sessions) - antirez#605: 1 commit (batched server session recovery race fix) - antirez#605 streaming improvements (4 commits): deferred — conflict with per-tier LRU, adapt later MTP/DSpark: present (from upstream fc9efd1 via feat/mtp-native). --batched-session --prefill-chunk 1024 now available (antirez#604).
iSevenDays
added a commit
to iSevenDays/ds4
that referenced
this pull request
Jul 25, 2026
…r LRU The integration commit 8577818 cherry-picked PR antirez#604 (session batching + VRAM reserve) onto feat/mtp-native, but the cherry-pick left 33 unresolved git conflict markers throughout ds4_cuda.cu, plus a duplicate definitions block. Each conflict was HEAD=per-tier LRU (our f836788 port, uses ssd_current() and g_ssd[t].expert_cache) vs 10ba298=upstream per-class (uses g_stream_expert_caches[DS4_CUDA_STREAM_EXPERT_CLASSES]). Resolution: - Resolved all 33 conflict regions in favor of HEAD (per-tier), dropping upstream's per-class infrastructure (g_stream_expert_caches[], g_stream_expert_runtime_caps[], g_stream_expert_class_*_bytes[], cuda_stream_expert_cache_class_index/_configured_budget_class/ _release_class, DS4_CUDA_STREAM_EXPERT_CLASSES define). - Removed the duplicate definitions block (the second enum, duplicate struct cuda_stream_expert_cache_slot/_cache, and the second g_stream_expert_budget_override) introduced by the cherry-pick; the original per-tier definitions at the top of the file are authoritative. - Dropped the leftover upstream direct-load block in cuda_stream_selected_cache_begin_load that referenced the undefined g_stream_selected_cache global and direct_loads/cache_hits/cache_misses counters (HEAD's per-tier flow already handles this via sc->gate_ptr and copied_from_global_cache). - Removed the duplicate no-op stub of cuda_model_load_progress_finish; the real implementation at line 1550 stands. - Kept cuda_stream_layer_expert_ranges_valid (pure bounds-check helper, no per-tier/per-class state) because it is still called from ds4_gpu_stream_expert_cache_seed_experts outside any conflict region. Build passes clean with CUDA_ARCH=sm_89. The resulting binary exposes all the expected flags: --ssd-streaming-cache-experts (LRU), --batched-session (PR antirez#604), --mtp*/--dspark* (MTP/DSpark). Co-Authored-By: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
--batched-session Nworks beautifully on multi-GPU placements, where thesession pipeline fills different pipeline stages across devices. On a single
GPU I ran into two things while testing it on an RTX PRO 6000 (96 GB) with
the Flash q2 GGUF, and this PR is an attempt at addressing both. Happy to
rework any part of it — the branch is organized so the pieces can be taken
(or dropped) independently.
1. Resident session creation fails at normal context sizes
On 96 GB,
--batched-session 4 --ctx 32768(and even--batched-session 1 --ctx 32768) fails withfailed to create cuda session: the lazy Q8→F16weight caches warm up during engine open and grow until only the default
reserve (~1% of VRAM) stays free, and the resident sessions are created
after that point. Measured on q2: engine base lands at ~95.9 GiB while each
32k session only needs ~0.6 GiB of context buffers.
Fix (2 commits): the server passes the resident session count to the engine
(
resident_sessionsconfig field), and engine open computes a cache reservefrom the per-session
ds4_context_memoryestimate right beforeaccelerator_cache_model_tensors()— the estimate needs the model shape, soit has to run inside open, not before. An explicit
DS4_CUDA_Q8_F16_CACHE_RESERVE_MBstill wins. With this,--batched-session 4 --ctx 32768 --prefill-chunk 1024boots at 97.2/97.9 GiB and single-streamspeed is unchanged (the hot part of the cache still fits).
2. Single-GPU decode batches don't scale
Without a placement,
ds4_sessions_eval_batch_cudaencodes the one-tokengraphs back to back. That preserves exactness, but each session re-reads the
full weight stream, so decode is memory-bound at 1x: measured 62 tok/s
aggregate at 2 and 4 concurrent clients vs 61 tok/s for one client.
This PR adds a fused single-tier path to the same function (the single-GPU
counterpart of
metal_graph_encode_session_pipeline_batch): the denseprojections run once on the shared batch workspace with
n_tokens = count,so the weights cross memory once per step; KV append and compressor/indexer
bookkeeping run per sequence with the same single-token helpers the decode
layer uses; attention runs as one multi-sequence launch with per-sequence
fallbacks (indexed path, score-buffer overflow), so there is no position
limit.
Numbers (RTX PRO 6000, Flash q2,
--batched-session 4 --ctx 32768 --prefill-chunk 1024, 200 generated tokens per request, aggregate =total tokens / wall):
Commit walkthrough
CUDA kernels (additive):
CUDA: add multi-sequence decode attention kernel— one query token persequence, each attending to its own KV ring/compressed cache
(
ds4_gpu_attn_seqview, by-value in kernel params, max 4).CUDA: add pos-vector RoPE tail variants— decode tokens of differentsequences sit at arbitrary positions; per-element math identical to the
single-token kernels (token t matches a pos0=pos[t] launch bitwise).
CUDA: narrow-batch q8_0 matmul for 2-4 token decode— same idea as yourtok2kernel, generalized to n_tok 2-4 with n_tok as a parameter: onewarp per output row, each 32-weight block read once and dotted against
all tokens. Dispatched ahead of cuBLAS/MMA for n_tok 2-4
(
DS4_CUDA_NO_Q8_NARROW=1restores the previous dispatch; the GLM verifyforce_decode_warppath is untouched).Graph driver:
metal_graph_decode_compressor_indexer— the per-token compressor/indexersegment of the decode layer phase, taken as a helper so the batched driver
can run it per sequence on batch-scratch row views. It is a copy (the
single-token path is untouched); if you'd rather have the phase function
call the helper too to avoid the duplication, glad to do that instead.
metal_graph_encode_layer_attention_decode_multi+metal_graph_encode_decode_multi— the fused pass. FFN and output headreuse the existing batch functions; logits land one row per sequence in
spec_logits.Graph: run fused-decode routed experts per sequence— at batch 2-4 theselected experts of different sequences are almost always disjoint, so the
batch expert path costs more than n single-token passes. This turned out
to matter twice: without it the fused path was slower than pipelining (54
vs 62) and one of two concurrent greedy streams diverged from the solo
run; with it, both concurrent greedy streams match the solo stream
token-for-token on the prompts I tested.
DS4_FUSED_BATCH_MOE=1keeps thebatch expert path for comparison.
Server/glue:
ds4_sessions_eval_batch_cudafor single-tier,non-streaming, non-distributed decode batches of 2-4;
DS4_NO_FUSED_SESSION_BATCH=1restores the one-token pipelining (and wasused for the A/B numbers above).
Correctness notes and limitations
reduction order); greedy token streams matched the solo run in every check
I ran (q2, 150-200 tokens, 2 and 4 concurrent). The kill switch makes the
comparison easy to reproduce, and I kept it default-on only because the
numbers looked convincing — happy to flip the default if you prefer
opt-in.
its current paths untouched.
the existing behavior).
spec_logitsis now allocated unless the fused path is disabled (~8 MiBper graph), since the fused output head lands there.
Test setup details, exact configs and raw runs are in the commit messages;
I can share the bench harness if useful. And if you had a different
direction in mind for single-GPU batching, happy to adapt this or just take
whatever pieces are useful.